package me.eg.fifth;

/**
 * 327. 区间和的个数
 * <p>
 * 给你一个整数数组 nums 以及两个整数 lower 和 upper 。
 * 求数组中，值位于范围 [lower, upper] （包含 lower 和 upper）之内的 区间和的个数 。
 * <p>
 * 区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤ j)。
 * <p>
 * 链接：https://leetcode-cn.com/problems/count-of-range-sum
 */
public class CountRangeSum {
    public int countRangeSum(int[] nums, int lower, int upper) {
        long[] preSum = new long[nums.length + 1];
        for (int i = 1; i <= nums.length; i++) {
            preSum[i] = preSum[i - 1] + nums[i - 1];
        }
        return countRangeSum(preSum, 0, preSum.length - 1, lower, upper);
    }

    private int countRangeSum(long[] preSum, int l, int r, int lower, int upper) {
        if (l >= r) return 0;
        int mid = (l + r) >> 1;
        int leftCnt = countRangeSum(preSum, l, mid, lower, upper);
        int rightCnt = countRangeSum(preSum, mid + 1, r, lower, upper);
        int cnt = countRange(preSum, l, r, mid, lower, upper);
        merge(preSum, l, r, mid);
        return leftCnt + rightCnt + cnt;
    }

    private int countRange(long[] preSum, int l, int r, int mid, int lower, int upper) {
        int ans = 0;
        int li = l, rs = mid + 1, re = mid + 1;
        while (li <= mid) {
            while (rs <= r && preSum[rs] - preSum[li] < lower) rs++;
            while (re <= r && preSum[re] - preSum[li] <= upper) re++;
            ans += re - rs;
            li++;
        }
        return ans;
    }

    private void merge(long[] preSum, int l, int r, int mid) {
        long[] temp = new long[r - l + 1];
        int left = l, right = mid + 1, index = 0;
        while (left <= mid || right <= r) {
            if (right > r || (left <= mid && preSum[left] <= preSum[right])) temp[index++] = preSum[left++];
            else temp[index++] = preSum[right++];
        }
        for (int i = 0; i < temp.length; i++) {
            preSum[l + i] = temp[i];
        }
    }
}
